#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
class KernelBitwiseLeftShift {
public:
    __aicore__ inline KernelBitwiseLeftShift() {}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR other, GM_ADDR out,
                                uint32_t smallCoreDataNum, uint32_t bigCoreDataNum, uint16_t bigCoreLoopNum, 
                                uint16_t smallCoreLoopNum, uint16_t ubPartDataNum, 
                                uint16_t smallCoreTailDataNum, uint16_t bigCoreTailDataNum, 
                                uint16_t tailBlockNum, TPipe* pipeIn){
        uint32_t globalBufferIndex = bigCoreDataNum * GetBlockIdx();
        this->ubPartDataNum = ubPartDataNum;
        if (GetBlockIdx() < tailBlockNum) 
        { 
            this->coreDataNum = bigCoreDataNum;
            this->tileNum = bigCoreLoopNum;
            this->tailDataNum = bigCoreTailDataNum;
        }
        else 
        { 
            this->coreDataNum = smallCoreDataNum;
            this->tileNum = smallCoreLoopNum;
            this->tailDataNum = smallCoreTailDataNum;
            globalBufferIndex -= (bigCoreDataNum - smallCoreDataNum) * (GetBlockIdx() - tailBlockNum);
        }
        inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + globalBufferIndex, this->coreDataNum);
        otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + globalBufferIndex, this->coreDataNum);
        outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + globalBufferIndex, this->coreDataNum);;
        if constexpr (std::is_same_v<DTYPE_INPUT, int8_t>)
        {
            pipe.InitBuffer(inQueueInput, BUFFER_NUM, this->ubPartDataNum * sizeof(DTYPE_INPUT));
            pipe.InitBuffer(inQueueOther, BUFFER_NUM, this->ubPartDataNum * sizeof(DTYPE_INPUT));
            pipe.InitBuffer(outQueueOut, BUFFER_NUM, this->ubPartDataNum * sizeof(DTYPE_INPUT));
            pipeIn->InitBuffer(QueueTmp1, this->ubPartDataNum * sizeof(half));
            pipeIn->InitBuffer(QueueTmp2, this->ubPartDataNum * sizeof(half));
        }
        else
        {
            pipe.InitBuffer(inQueueOther, BUFFER_NUM, 65280); 
            pipe.InitBuffer(inQueueInput, BUFFER_NUM, 65792); 
            pipe.InitBuffer(outQueueOut, BUFFER_NUM, this->ubPartDataNum * sizeof(DTYPE_INPUT));
        }
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->tileNum-1;
        this->processDataLenght = this->ubPartDataNum;
        for (int32_t i = 0; i < loopCount; i++) 
        {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataLenght = this->tailDataNum;
        CopyIn(loopCount);
        Compute(loopCount);
        CopyOut(loopCount);
    }
private:
 __aicore__ inline void CopyIn(int32_t progress) {
    LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.AllocTensor<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.AllocTensor<DTYPE_INPUT>();
        DataCopy(inputLocal, inputGm[progress * this->ubPartDataNum], this->processDataLenght);
        DataCopy(otherLocal, otherGm[progress * this->ubPartDataNum], this->processDataLenght);
        inQueueInput.EnQue(inputLocal);
        inQueueOther.EnQue(otherLocal);
    }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.AllocTensor<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, int8_t>)
        {
            auto tmp1 = QueueTmp1.Get<half>();
            auto tmp2 = QueueTmp2.Get<half>();
            Cast(tmp1, inputLocal, RoundMode::CAST_NONE, this->processDataLenght);
            Cast(tmp2, otherLocal, RoundMode::CAST_NONE, this->processDataLenght);
            Muls(tmp2, tmp2, (half)0.69314718055, this->processDataLenght);
            Exp(tmp2, tmp2, this->processDataLenght);
            MulCast(outLocal, tmp1, tmp2, this->processDataLenght);
        }
        else 
        {
            LocalTensor<int32_t> otherLocal_int32 = otherLocal.ReinterpretCast<int32_t>();
            LocalTensor<uint32_t> otherLocal_uint32 = otherLocal.ReinterpretCast<uint32_t>();
            LocalTensor<int32_t> inputLocal_int32 = inputLocal.ReinterpretCast<int32_t>();
            Adds(otherLocal_int32, otherLocal_int32, (int32_t)127, this->processDataLenght);
            ShiftLeft(otherLocal_uint32,otherLocal_uint32, (uint32_t)23, this->processDataLenght);
            Cast(otherLocal_int32, otherLocal.ReinterpretCast<float>(), RoundMode::CAST_RINT, this->processDataLenght);
            Mul(outLocal.ReinterpretCast<int32_t>(), inputLocal_int32, otherLocal_int32, this->processDataLenght);
        }
        inQueueInput.FreeTensor(inputLocal);
        inQueueOther.FreeTensor(otherLocal);
        outQueueOut.EnQue<DTYPE_INPUT>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.DeQue<DTYPE_INPUT>();
        DataCopy(outGm[progress * this->ubPartDataNum], outLocal, this->processDataLenght);
        outQueueOut.FreeTensor(outLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, 1> inQueueInput, inQueueOther;
    TQue<QuePosition::VECOUT, 1> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;
    GlobalTensor<DTYPE_INPUT> inputGm;
    GlobalTensor<DTYPE_INPUT> otherGm;
    GlobalTensor<DTYPE_INPUT> outGm;
    uint16_t processDataLenght;
    uint32_t coreDataNum;
    uint16_t tileNum;
    uint16_t ubPartDataNum;
    uint16_t tailDataNum;
};
class KernelBitwiseLeftShiftBC_1 {
public:
    __aicore__ inline KernelBitwiseLeftShiftBC_1() {}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR other, GM_ADDR out,
                                uint32_t bigNum, uint32_t bigLength, uint32_t smallLength, 
                                uint32_t repeat, uint32_t cycles, uint32_t interval, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreDataNum;
        uint32_t group = GetBlockIdx()/repeat;
        uint32_t within = GetBlockIdx()%repeat;
        if (within < bigNum) {
            coreDataNum = bigLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + group * cycles * interval + within * bigLength * interval);
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + group * interval);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + group * cycles * interval + within * bigLength * interval);
        } else {
            coreDataNum = smallLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + group * cycles * interval + bigLength * interval * bigNum + smallLength * interval  * (within - bigNum));
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + group * interval);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + group * cycles * interval + bigLength * interval * bigNum + smallLength * interval  * (within - bigNum));
        }
        this->cycles = cycles;
        this->interval = interval;
        if constexpr (std::is_same_v<DTYPE_INPUT, int8_t>)
        {
            this->tileDataLenght = 9824;
        }
        else if constexpr (std::is_same_v<DTYPE_INPUT, int16_t>)
        {
            this->tileDataLenght = 4896;
        }
        else if constexpr (std::is_same_v<DTYPE_INPUT, int32_t>)
        {
            this->tileDataLenght = 8192;
        }
        else
        {
            this->tileDataLenght = 8192;
        }
        this->extendNum = 32;
        this->tileDataLenght = this->extendNum * interval;
        this->loopCount = coreDataNum/this->extendNum;
        this->tailDataLenght = (coreDataNum - this->extendNum*this->loopCount) * interval;
        if(this->tailDataLenght == 0){this->tailDataLenght = this->tileDataLenght; this->loopCount--;}
        pipeIn->InitBuffer(inQueueInput, BUFFER_NUM, this->tileDataLenght  * sizeof(DTYPE_INPUT));
        pipeIn->InitBuffer(inQueueOther, BUFFER_NUM, this->tileDataLenght  * sizeof(DTYPE_INPUT));
        pipeIn->InitBuffer(outQueueOut, BUFFER_NUM, this->tileDataLenght  * sizeof(DTYPE_INPUT));
    }
    __aicore__ inline void Process() {
        this->processDataLenght = this->tileDataLenght;
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.AllocTensor<DTYPE_INPUT>();
        DataCopy(otherLocal, otherGm, this->interval);
        inQueueOther.EnQue(otherLocal);
        otherLocal = inQueueOther.DeQue<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, int64_t>) {
            Cast(otherLocal.ReinterpretCast<float>(), otherLocal, RoundMode::CAST_RINT, this->interval);
            Muls(otherLocal.ReinterpretCast<float>(), otherLocal.ReinterpretCast<float>(), (float)0.69314718055994530941723212145818, this->interval);
            Exp(otherLocal.ReinterpretCast<float>(), otherLocal.ReinterpretCast<float>(), this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[2*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 2*this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[4*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 4*this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[8*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 8*this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[16*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 16*this->interval);
        }
        for (int32_t i = 1; i < this->loopCount; i++) {
            CopyIn(i);
            Compute(i, otherLocal);
            CopyOut(i);
        }
        this->processDataLenght = this->tailDataLenght;
        CopyIn(loopCount);
        Compute(loopCount, otherLocal);
        CopyOut(loopCount);
        inQueueOther.FreeTensor(otherLocal);
    }
private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.AllocTensor<DTYPE_INPUT>();
        DataCopy(inputLocal, inputGm[progress * this->tileDataLenght], this->processDataLenght);
        inQueueInput.EnQue(inputLocal);
    }
    __aicore__ inline void Compute(int32_t progress, LocalTensor<DTYPE_INPUT> otherLocal) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.AllocTensor<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, int64_t>) {
            Cast(inputLocal.ReinterpretCast<float>(), inputLocal, RoundMode::CAST_RINT, this->processDataLenght);
            Mul(inputLocal.ReinterpretCast<float>(), inputLocal.ReinterpretCast<float>(), otherLocal.ReinterpretCast<float>(), this->processDataLenght);
            Cast(outLocal, inputLocal.ReinterpretCast<float>(), RoundMode::CAST_RINT, this->processDataLenght);
        }
        inQueueInput.FreeTensor(inputLocal);
        outQueueOut.EnQue<DTYPE_INPUT>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.DeQue<DTYPE_INPUT>();
        DataCopy(outGm[progress * this->tileDataLenght], outLocal, this->processDataLenght);
        outQueueOut.FreeTensor(outLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueInput, inQueueOther;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,QueueTmp3,QueueTmp4, QueueTmp5;
    GlobalTensor<DTYPE_INPUT> inputGm;
    GlobalTensor<DTYPE_INPUT> otherGm;
    GlobalTensor<DTYPE_INPUT> outGm;
    uint32_t cycles;
    uint32_t interval;
    uint32_t loopCount;
    uint32_t tileDataLenght;
    uint32_t tailDataLenght;
    uint32_t processDataLenght;
    uint32_t extendNum;
};
extern "C" __global__ __aicore__ void bitwise_left_shift(GM_ADDR input, GM_ADDR other, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    TPipe pipe;
    if(TILING_KEY_IS(1)) {
        GET_TILING_DATA_WITH_STRUCT(BitwiseLeftShiftTilingData, tiling_data, tiling);
        KernelBitwiseLeftShift op;
        op.Init(input, other, out,tiling_data.smallCoreDataNum,                                
                tiling_data.bigCoreDataNum, tiling_data.bigCoreLoopNum,             
                tiling_data.smallCoreLoopNum, tiling_data.ubPartDataNum,           
                tiling_data.smallCoreTailDataNum, tiling_data.bigCoreTailDataNum,   
                tiling_data.tailBlockNum, &pipe);
        op.Process();
    }
    else if(TILING_KEY_IS(3)) {
        GET_TILING_DATA_WITH_STRUCT(BitwiseLeftShiftTilingDataBC, tiling_data, tiling);
        KernelBitwiseLeftShiftBC_1 op;
        op.Init(input, other, out,
                tiling_data.bigNum, tiling_data.bigLength, tiling_data.smallLength, 
                tiling_data.repeat, tiling_data.cycles, tiling_data.interval, &pipe);
        op.Process();
    }
}